import numpy as np
from scipy.optimize import curve_fit

def gaussian_2d(x,y,x_mean,y_mean,sigma,amp,bias:0):
    r2 = (x-x_mean)**2+(y-y_mean)**2
    return amp*np.exp(-r2/(2*sigma**2)) + bias

def gaussian(x, mu, sigma, amp, bias):
    return amp*np.exp(-(x-mu)**2/(2*sigma**2)) + bias

def calibrate_emccd_fit(count_matrix):
    y = np.sum(count_matrix, axis=0)
    x = range(len(y))
    bounds = (np.array([0, 0, 0, -np.inf]), [np.inf, np.inf, np.inf, np.inf])
    params_cali, cov = curve_fit(gaussian, x, y, p0=[np.median(x), len(x)/2, 100, 0], bounds=bounds)
    
    x_fit = np.arange(x[0],x[-1],(x[-1]-x[0])/50)
    plt.figure()
    plt.plot(x, count_matrix_1d, 'o')
    plt.plot(x_fit, gaussian(x_fit, params_cali[0], params_cali[1], params_cali[2], params_cali[3]))
    plt.show()
    return {'mu': params_cali[0], 'sigma': params_cali[1], 'amp': params_cali[2], 'bias': params_cali[3]}

def emccd_fit_once(x, count_matrix, params_cali):
    y = np.sum(count_matrix, axis=0)
    bounds = (np.array([params_cali['mu'], params_cali['sigma'], 0, params_cali['bias']])-1e-5, [params_cali['mu'], params_cali['sigma'], np.inf, params_cali['bias']])
    params, cov = curve_fit(gaussian, x, y, bounds=bounds)
    return params[2]

def get_infidelity(amp_list_dark, amp_list_bright, bins=range(1000)):
    p, x = np.histogram(amp_list_dark, bins=bins)
    p = p/sum(p)
    dark_dis = [p, x]
    
    p, x = np.histogram(amp_list_bright, bins=bins)
    p = p/sum(p)
    bright_dis = [p, x]
    threshold_list, infidelity_list = [dark_dis[1][:-1], [(sum(dark_dis[0][i:]) + sum(bright_dis[0][:i]))/2 for i in range(len(dark_dis[0]))]]
    infidelity = min(infidelity_list)
    threshold  = threshold_list[infidelity_list.index(infidelity)]
    return [threshold_list, infidelity_list, threshold, infidelity]